from P03_NER.LSTM_CRF.config import Config

conf = Config()
# print(conf.train_path)

# 构造数据集：对train.txt进行处理，按照句子结束标点进行切分，得到x,y样本对，放到列表中
def build_data():
    datas = []  # 最终存储数据的三维列表
    sample_x = []  # 存储一个样本的文字数据
    sample_y = []  # 存储一个样本的标签数据
    vocab_list = ['PAD', 'UNK']  # 词汇列表，初始化两个特殊字符

    # 遍历train.txt数据，将数据存到列表中
    with open(conf.train_path, 'r', encoding='utf-8') as fr:
        for line in fr.readlines():
            # print(f'line-->{line}')
            word_tag_list = line.strip('\n').split('\t')  # 为了防止把数据中的空行处理掉，这里不再直接使用strip()
            # 判断是否为空行
            # if not word_tag_list: # 如果是空行，则跳过
            #     continue
            if len(word_tag_list) != 2:
                # print(f'line-->{line}')
                continue

            word = word_tag_list[0]
            if not word: # 如果word为None或''，则跳过
                continue
            tag = word_tag_list[-1]

            # 将word和tag添加到列表中
            sample_x.append(word)
            sample_y.append(tag)
            # print(f'sample_x--->{sample_x}')
            # print(f'sample_y--->{sample_y}')

            # 如果word在结束的标点符号中，则认为句子结束了，将sample_x和sample_y组装成一个列表后添加到datas中
            if word in ['。', '?', '!', '！', '？']:
                datas.append([sample_x, sample_y])
                # print(f'datas--->{datas}')
                # 清空sample_x和sample_y
                sample_x = []
                sample_y = []

            # 将没有重复的文字添加到词汇列表中
            if word not in vocab_list:
                vocab_list.append(word)
        # 为了方便后续使用这个列表，将其保存到文件中
        with open(conf.vocab_path, 'w', encoding='utf-8') as fw:
            fw.write('\n'.join(vocab_list))
        # 构造word2id字典
        word2id = {word: i for i, word in enumerate(vocab_list)}

        # 返回datas和word2id
        return datas, word2id


if __name__ == '__main__':
    datas, word2id = build_data()
    # print(f'datas-->{datas}')
    # print(f'word2id-->{word2id}')
    print(f'len(datas)-->{len(datas)}')  # 7914
    print(f'len(word2id)-->{len(word2id)}')  # 1759